import re
import pandas as pd
from nltk.corpus import stopwords
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

stop_words = set(stopwords.words("english"))


class TopicMiningPro:

    # 文本清理函数
    def preprocess_text(self, text):
        if pd.isna(text):  # 检查是否为 NaN
            return ""  # 返回空值作为默认值
        # 小写化
        text = text.lower()
        # 去除特殊字符和数字
        text = re.sub(r"[^a-zA-Z\s]", "", text)
        # 去掉停用词
        text = " ".join(word for word in text.split() if word not in stop_words)
        return text

    def __init__(self, filename):
        # 读取数据
        self.data = pd.read_csv("../InputData/" + filename + ".tsv", sep="\t")

        # 删除缺失值
        self.data.dropna(
            subset=["benefitsReview", "sideEffectsReview", "commentsReview"],
            inplace=True,
        )

        # 应用预处理
        self.data["benefitsReview_cleaned"] = self.data["benefitsReview"].apply(
            self.preprocess_text
        )
        self.data["sideEffectsReview_cleaned"] = self.data["sideEffectsReview"].apply(
            self.preprocess_text
        )
        self.data["commentsReview_cleaned"] = self.data["commentsReview"].apply(
            self.preprocess_text
        )

        # 合并文本字段
        self.data["combinedText"] = (
            self.data["benefitsReview_cleaned"]
            + " "
            + self.data["sideEffectsReview_cleaned"]
            + " "
            + self.data["commentsReview_cleaned"]
        )

        # 将文本和药品名称分开
        self.X = self.data["combinedText"]
        self.y = self.data["urlDrugName"]

    def train_model(self):
        # 划分训练和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=0.2, random_state=42
        )

        # 使用TF-IDF把文本特征向量化为数值
        tfidf = TfidfVectorizer(max_features=10000)
        X_train_tfidf = tfidf.fit_transform(X_train)
        X_test_tfidf = tfidf.transform(X_test)

        # 训练一个逻辑回归分类器
        model = LogisticRegression(max_iter=1000)
        model.fit(X_train_tfidf, y_train)

        # 在测试集上评估
        y_pred = model.predict(X_test_tfidf)
        print("Test Accuracy:", round(accuracy_score(y_test, y_pred), 2))


if __name__ == "__main__":
    example = TopicMiningPro("drugLibTrain_raw")
    example.train_model()
